#include "sci_lio/voxel_map.h"

//根据配置参数初始化一个体素地图
VoxelMap::VoxelMap(Config config) {
  //通过配置参数设置体素的分辨率，影响体素的大小
  resolution_ = config.resolution;
  //计算分辨率的倒数，用于后续的体素索引计算
  inv_resolution_ = 1.0 / resolution_;
  //定义体素映射的容量，即所能存储的最大体素数量
  capacity_ = config.capacity;
  //每个体素中最多可以存储的点数
  grid_max_points_ = config.grid_max_points;

  //初始化 delta_P_，用于存储查找相邻体素时需要的偏移向量
  delta_P_.clear();

  voxel_MaxZ_map_ = std::make_shared<MyHashMapVoxel>();
  //代码根据 config.search_method 的不同值来填充 delta_P_ 向量，确定检查邻近体素的方式
  //对于 "NEARBY_1"，只添加了中心点
  if (config.search_method == "NEARBY_1") {
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, 0.0));
  } 
  //! "NEARBY_7" 添加了中心点及其在六个方向的相邻点（前、后、左、右、上、下）
  else if (config.search_method == "NEARBY_7") {
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, resolution_));
  }
  //"NEARBY_19" 和 "NEARBY_26" 分别添加了更多的周围点，通过组合不同的分辨率偏移量以覆盖更广的邻域 
  else if (config.search_method == "NEARBY_19") {
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, -resolution_));
  } else if (config.search_method == "NEARBY_26") {
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, -resolution_, 0.0));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, 0.0, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(0.0, -resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, -resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(resolution_, -resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, resolution_, -resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, -resolution_, resolution_));
    delta_P_.push_back(Eigen::Vector3d(-resolution_, -resolution_, -resolution_));
  }
  //为临时存储体素数据的哈希映射和向量分配内存，使用智能指针进行管理，提高内存安全性
  temp_voxel_map_ptr_ = std::make_shared<MyHashMap>();
  temp_voxel_array_ptr_ = std::make_shared<MyVector>();

  search_range_.clear();
  for (int x_gain = -1; x_gain <= 1; ++x_gain) {
    for (int y_gain = -1; y_gain <= 1; ++y_gain) {
      for (int z_gain = -1; z_gain <= 1; ++z_gain) {
        search_range_.emplace_back(Eigen::Vector3d(x_gain * resolution_,
                                                    y_gain * resolution_,
                                                    z_gain * resolution_));
      }
    }
  }

}

//将输入的点云数据添加到体素地图中，并处理相应的数据以维护体素的统计信息
/*
检查并处理点云数据。
通过哈希索引管理体素的创建和更新。
计算并存储每个体素的统计信息，包括质心、协方差和点数量。
维护体素地图的大小，使用 LRU 缓存机制优化内存管理。
*/
bool VoxelMap::AddCloud(const CloudPtr& input_cloud_ptr) //input_cloud_ptr是矫正后的，转换到世界坐标系下的点云
{
  //如果输入点云是空的，就返回
  if (input_cloud_ptr->empty()) 
  {
    LOG(INFO) << "input cloud is empty";
    return false;
  }

  //创建一个向量 point_buff，用于存储所有点的哈希索引及其坐标
  std::vector<point_hash_idx> point_buff;
  point_buff.resize(input_cloud_ptr->size());
  //使用 Intel TBB 库对输入点云中的每个点并行计算其哈希索引。这段代码会为每个点生成一个 point_hash_idx 结构体，并将计算得到的点和哈希索引存储在 point_buff 中
  tbb::parallel_for(tbb::blocked_range<size_t>(0, input_cloud_ptr->size()),[&, this](tbb::blocked_range<size_t> r) 
  {
    for (size_t i = r.begin(); i < r.end(); ++i) 
    {
      point_hash_idx p;
      p.point_ = input_cloud_ptr->points[i].getVector3fMap().cast<double>();
      p.hash_idx_ = ComputeHashIndex(p.point_);
      p.hash_idx_z_ = ComputeMaxZHashIndex(p.point_);
      point_buff[i] = p;
    }
  });

  for (size_t i = 0; i < point_buff.size();i++) 
  {
    size_t curr_hash_idx_z = point_buff.at(i).hash_idx_z_;

    MyAccessorVoxel Z_accessor;
    voxel_MaxZ_map_->insert(Z_accessor, curr_hash_idx_z);
    if (Z_accessor->second == nullptr) 
    {
      Z_accessor->second = std::make_shared<Voxel>();
      Z_accessor->second->max_z = point_buff.at(i).point_[2];
      Z_accessor->second->min_z = point_buff.at(i).point_[2];
      Z_accessor->second->maxZ_hash_index_ =  point_buff.at(i).hash_idx_;
      Z_accessor->second->minZ_hash_index_ =  point_buff.at(i).hash_idx_;
    }
    else 
    {
      if(point_buff.at(i).point_[2]>Z_accessor->second->max_z)
      {
        Z_accessor->second->max_z = point_buff.at(i).point_[2];
        Z_accessor->second->maxZ_hash_index_ = point_buff.at(i).hash_idx_;
      }
      if(point_buff.at(i).point_[2]<Z_accessor->second->min_z)
      {
        Z_accessor->second->min_z = point_buff.at(i).point_[2];
        Z_accessor->second->minZ_hash_index_ = point_buff.at(i).hash_idx_;
      }
    }
  }
  
  //对点缓冲区进行排序，以便后续处理时能快速查找相同哈希索引的点
  tbb::parallel_sort(point_buff.begin(), point_buff.end());

  size_t error_grids = 0;
  //! 外层循环迭代 point_buff，处理具有相同哈希索引的点。这部分代码中主要进行统计和重建体素的过程。
  for (size_t i = 0; i < point_buff.size();) 
  {
    size_t j = i;
    size_t curr_hash_idx = point_buff.at(i).hash_idx_;
    Eigen::Vector3d point_sum = Eigen::Vector3d::Zero();
    Eigen::Matrix3d cov_sum = Eigen::Matrix3d::Zero();
    size_t count = 0;
    std::vector<Eigen::Vector3d> points_array_temp;
    //在内部循环中，检查当前哈希索引是否已经存在于体素地图中。根据情况创建新体素或更新现有体素的统计信息
    for (; j < point_buff.size(); ++j) 
    {
      //! 如果下一个点和当前点的hash索引一致，表示他们在一个体素中，则计算位置和（point_sum），以及位置平方的和cov_sum，并count+1
      if (curr_hash_idx == point_buff.at(j).hash_idx_) 
      {
        // prevent hash collision
        if (IsSameGrid(point_buff.at(i).point_, point_buff.at(j).point_)) 
        {
          point_sum += point_buff.at(j).point_;
          cov_sum += point_buff.at(j).point_ * point_buff.at(j).point_.transpose();
          count++;
          points_array_temp.emplace_back(point_buff.at(j).point_);
        } 
        else 
          error_grids++; // hash collision!
      } 
      else
        break;
    }

    auto iter = voxel_map_.find(curr_hash_idx);
    //如果地图中没有这个体素，创建新的体素并计算其统计信息（包括质心和协方差）
    if (iter == voxel_map_.end()) 
    {
      std::shared_ptr<Grid> grid = std::make_shared<Grid>(grid_max_points_);
      grid->points_sum_ = point_sum;
      grid->points_num_ = count;
      grid->cov_sum_ = cov_sum;
      grid->points_array_ = points_array_temp;
      //计算质心
      grid->centroid_ = grid->points_sum_ / static_cast<double>(grid->points_num_);
      //计算协方差矩阵
      ComputeCovariance(grid);

      //维护一个 LRU (Least Recently Used) 缓存，确保体素地图的大小不超过设定的最大容量
      grids_cache_.push_front({curr_hash_idx, grid});
      voxel_map_.insert({curr_hash_idx, grids_cache_.begin()});
      if (voxel_map_.size() >= capacity_) 
      {
        voxel_map_.erase(grids_cache_.back().first);
        grids_cache_.pop_back();
      }
    } 
    //如果哈希索引存在，则根据点数量、是否发生哈希碰撞等条件更新已有体素的信息,包括质心、协方差和点数量
    else
    {
      Eigen::Vector3d centroid = point_sum / static_cast<double>(count);
      // prevent hash collision
      if (!IsSameGrid(iter->second->second->centroid_, centroid)) 
      {
        // hash collision!
        error_grids++;
        // remove this grid
        grids_cache_.erase(iter->second);

        // create a new grid
        std::shared_ptr<Grid> grid = std::make_shared<Grid>(grid_max_points_);
        grid->points_sum_ = point_sum;
        grid->points_num_ = count;
        grid->cov_sum_ = cov_sum;
        grid->points_array_ = points_array_temp;
        // compute centroid
        grid->centroid_ = grid->points_sum_ / static_cast<double>(grid->points_num_);
        // compute covariance
        ComputeCovariance(grid);

        grids_cache_.push_front({curr_hash_idx, grid});
        voxel_map_[curr_hash_idx] = grids_cache_.begin();
      } 
      else 
      {
        //TODO::地图中的体素应该是动态更新的啊！ 只有在体素中点的数量少于50时，才会进行更新，确保数据的稳定性。
        if (iter->second->second->points_num_ < 50) 
        {
          iter->second->second->points_sum_ += point_sum;
          iter->second->second->points_num_ += count;
          iter->second->second->cov_sum_ += cov_sum;
          //计算质心
          iter->second->second->centroid_ = iter->second->second->points_sum_ / static_cast<double>(iter->second->second->points_num_);
          //计算协方差阵
          ComputeCovariance(iter->second->second);

          //限制每个体素中点的数量，来提升KNN的效率Improve KNN efficiency by limiting the number of points per
          // grid.
          if (iter->second->second->points_num_ < grid_max_points_) 
          {
            iter->second->second->points_array_.insert(
                iter->second->second->points_array_.end(),
                points_array_temp.begin(),
                points_array_temp.end());
          }
        }
        grids_cache_.splice(grids_cache_.begin(), grids_cache_, iter->second);
        voxel_map_[curr_hash_idx] = grids_cache_.begin();
      }
    }
    i = j;
  }

  LOG(INFO) << "加入关键帧：total_point " << input_cloud_ptr->size() << " error_point: " << error_grids;
  return true;
}

//计算协方差矩阵
void VoxelMap::ComputeCovariance(std::shared_ptr<Grid>& grid_ptr) 
{
  if (grid_ptr->points_num_ >= 6) 
  {
    Eigen::Matrix3d covariance =(grid_ptr->cov_sum_ -grid_ptr->points_sum_ * grid_ptr->centroid_.transpose()) /(static_cast<double>(grid_ptr->points_num_) - 1.0);

    Eigen::JacobiSVD<Eigen::Matrix3d> svd(covariance, Eigen::ComputeFullU | Eigen::ComputeFullV);

    Eigen::Vector3d values(1, 1, 1e-3);
    Eigen::Matrix3d modified_cov = svd.matrixU() * values.asDiagonal() * svd.matrixV().transpose();

    grid_ptr->inv_cov_ = modified_cov.inverse();
    grid_ptr->cov_ = modified_cov;
    grid_ptr->is_valid_ = true;
  } 
  else
  {
    // The number of points is too little to calculate a valid probability.
    grid_ptr->is_valid_ = false;
  }
}

//通过给定的条件查找在指定范围内最近的 K 个点
/*
根据给定的目标点、点数要求 K 和范围，从体素地图中查找在该范围内的最近 K 个点。
它通过计算每个邻近点与目标点之间的距离，筛选出符合条件的点，并对结果进行排序，最终返回符合条件的点列表
*/
bool VoxelMap::KNNByCondition(const Eigen::Vector3d& point,
                              const size_t K,
                              const double range,
                              std::vector<Eigen::Vector3d>& results) 
{
  //point_dist 向量来存储满足条件的点的信息
  std::vector<point_distance> point_dist;
  point_dist.reserve(delta_P_.size() * grid_max_points_);
  //range2 存储范围的平方，以避免在后面的比较中重复计算平方根
  double range2 = range * range;

  //遍历 delta_P_ 向量中的每个偏移量（即周围体素的相对位置），计算与目标点相邻的点 nearby_point
  for (const auto& delta : delta_P_) 
  {
    Eigen::Vector3d nearby_point = point + delta;
    //使用 ComputeHashIndex 函数计算邻近点的哈希索引
    size_t hash_idx = ComputeHashIndex(nearby_point);
    Eigen::Vector3d centroid;
    //在 voxel_map_ 中查找计算得到的哈希索引以获取对应的体素信息
    auto iter = voxel_map_.find(hash_idx);
    //使用 IsSameGrid 函数验证 nearby_point 是否确实属于该体素
    if (iter != voxel_map_.end() && IsSameGrid(nearby_point, iter->second->second->centroid_)) 
    {
      //如果体素存在，遍历该体素所有的点，计算其与目标点之间的距离（采用平方距离以避免开方计算）
      for (const auto& p : iter->second->second->points_array_) 
      {
        double dist = (point - p).squaredNorm();

        //如果距离在指定范围内，将其添加到 point_dist 向量中。
        if (dist < range2) 
        {
          point_dist.emplace_back(point_distance(
              p, dist, (&p - iter->second->second->points_array_.data())));
        }
      }
    }
   }

  if (point_dist.empty()) 
  {
    return false;
  }

  //如果找到的点数超过 K，使用 std::nth_element 只保留 K 个距离最近的点
  if (point_dist.size() > K) 
  {
    std::nth_element(point_dist.begin(), point_dist.begin() + K - 1, point_dist.end());
    point_dist.resize(K);
  }
  std::nth_element(point_dist.begin(), point_dist.begin(), point_dist.end());

  //清空 results 向量，将找到的最近点添加到 results 中
  results.clear();
  for (const auto& it : point_dist) {results.emplace_back(it.point_);}
  return true;
}

bool VoxelMap::GetCentroidAndCovariance(const size_t hash_idx,Eigen::Vector3d& centorid,Eigen::Matrix3d& cov) 
{
  auto iter = voxel_map_.find(hash_idx);
  if (iter != voxel_map_.end() && iter->second->second->is_valid_) 
  {
    centorid = iter->second->second->centroid_;
    cov = iter->second->second->cov_;
    return true;
  } 
    return false;
}

bool VoxelMap::GetCentroidAndCovariance(const size_t hash_idx,Eigen::Vector3d& centorid,Eigen::Matrix3d& cov,
                                        std::bitset<27>& features,bool& is_valid) 
{
  auto iter = voxel_map_.find(hash_idx);
  //能找到，但是这个栅格中的点数量小于6，导致cov计算不合格啊，也用不了
  if(calculate_invalid && iter != voxel_map_.end() && !iter->second->second->is_valid_)
    invalid_point1++;
  //找不到hash值，那确实没有办法啊
  if(calculate_invalid && iter == voxel_map_.end())
    invalid_point2++;

  if (iter != voxel_map_.end()) 
  {
    centorid = iter->second->second->centroid_;
    is_valid = iter->second->second->is_valid_;
    cov = iter->second->second->cov_;

    //计算二值特征
    for (size_t j = 0; j < search_range_.size(); ++j) 
    {
      Eigen::Vector3d near_point = iter->second->second->centroid_ + search_range_[j];
      size_t hash_idx_ = ComputeHashIndex(near_point);

      auto iter_ = voxel_map_.find(hash_idx_);
      if (iter_ != voxel_map_.end()) 
      {
        features.set(j);
      }
      else
      {
        features.reset(j);
      }
      //std::cout << "features: " << features << std::endl;
      //features = iter->second->second->features;

    }
    return true;
  } 
  else 
  {
    return false;
  }
}

size_t VoxelMap::ComputeHashIndex(const Eigen::Vector3d& point) {
  double loc_xyz[3];
  for (size_t i = 0; i < 3; ++i) {
    loc_xyz[i] = point[i] * inv_resolution_;
    if (loc_xyz[i] < 0) {
      loc_xyz[i] -= 1.0;
    }
  }

  size_t x = static_cast<size_t>(loc_xyz[0]);
  size_t y = static_cast<size_t>(loc_xyz[1]);
  size_t z = static_cast<size_t>(loc_xyz[2]);

  return ((((z)*HASH_P_) % MAX_N_ + (y)) * HASH_P_) % MAX_N_ + (x);
}

size_t VoxelMap::ComputeMaxZHashIndex(const Eigen::Vector3d& point) {
  double loc_xyz[3];
  for (size_t i = 0; i < 3; ++i) 
  {
    loc_xyz[i] = point[i] * inv_resolution_;
    if (loc_xyz[i] < 0) {
      loc_xyz[i] -= 1.0;
    }
  }

  size_t x = static_cast<size_t>(loc_xyz[0]);
  size_t y = static_cast<size_t>(loc_xyz[1]);
  size_t z = 0;

  return ((((z)*HASH_P_) % MAX_N_ + (y)) * HASH_P_) % MAX_N_ + (x);
}


bool VoxelMap::IsSameGrid(const Eigen::Vector3d& p1,
                          const Eigen::Vector3d& p2) {
  int hx_1 = floor(p1.x() * inv_resolution_);
  int hy_1 = floor(p1.y() * inv_resolution_);
  int hz_1 = floor(p1.z() * inv_resolution_);
  int hx_2 = floor(p2.x() * inv_resolution_);
  int hy_2 = floor(p2.y() * inv_resolution_);
  int hz_2 = floor(p2.z() * inv_resolution_);

  return ((hx_1 == hx_2) && (hy_1 == hy_2) && (hz_1 == hz_2));
}